import cv2
import numpy as np
import os
import json
from ultralytics import YOLO
from PIL import Image, ImageDraw, ImageFont
import numpy as np

class KeypointAnnotator:
    def __init__(self, image_dir):
        self.image_dir = image_dir
        self.images = [f for f in os.listdir(image_dir) if f.endswith(('.jpg', '.png'))]
        self.current_image_idx = 0
        self.keypoints = []
        self.current_point = 0
        self.window_name = "Keypoint Annotation Tool"
        self.keypoint_names = [
            "头部", "左肩", "右肩", 
            "左肘", "右肘", "左手腕", "右手腕",
            "左臀", "右臀", "左膝", "右膝",
            "左脚踝", "右脚踝"
        ]
        self.model = YOLO('yolov8n-pose.pt')
        self.current_image = None
        self.image_height = 0
        self.image_width = 0
        self.all_detected_keypoints = []
        self.selected_person = 0
        
        # 设置中文字体
        font_path = "SimHei.ttf"  # 使用已有的黑体字体
        self.font_size = 30
        self.font = ImageFont.truetype(font_path, self.font_size)
        self.small_font = ImageFont.truetype(font_path, 20)
        
    def put_chinese_text(self, img, text, position, font=None, color=(0, 255, 0)):
        """使用PIL添加中文文字到图片上"""
        if font is None:
            font = self.font
            
        # 将OpenCV图片转换为PIL图片
        pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        
        # 创建绘图对象
        draw = ImageDraw.Draw(pil_img)
        
        # 添加文字
        draw.text(position, text, font=font, fill=color)
        
        # 转换回OpenCV格式
        return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)

    def load_image(self):
        """加载当前图片并进行预标注"""
        image_path = os.path.join(self.image_dir, self.images[self.current_image_idx])
        self.current_image = cv2.imread(image_path)
        self.image_height, self.image_width = self.current_image.shape[:2]
        
        # 使用YOLOv8进行预标注
        results = self.model(self.current_image)
        self.all_detected_keypoints = []
        
        if len(results[0].keypoints.data) > 0:
            # 存储所有检测到的人的关键点
            for person_kpts in results[0].keypoints.data:
                person_points = []
                for kpt in person_kpts:
                    x, y = int(kpt[0]), int(kpt[1])
                    if x >= 0 and y >= 0:  # 只添加有效的关键点
                        person_points.append((x, y))
                    else:
                        person_points.append(None)
                self.all_detected_keypoints.append(person_points)
            
            # 如果还没有选择关键点，使用第一个检测到的人
            if not self.keypoints:
                self.select_person(0)
    
    def select_person(self, person_idx):
        """选择要标注的人"""
        if 0 <= person_idx < len(self.all_detected_keypoints):
            self.selected_person = person_idx
            self.keypoints = self.all_detected_keypoints[person_idx].copy()
            self.current_point = len([p for p in self.keypoints if p is not None])
        
    def mouse_callback(self, event, x, y, flags, param):
        if event == cv2.EVENT_LBUTTONDOWN:
            if self.current_point < len(self.keypoint_names):
                self.keypoints[self.current_point] = (x, y)
                self.current_point += 1
                self.update_image()
                
    def update_image(self):
        if self.current_image is None:
            self.load_image()
            
        display_image = self.current_image.copy()
        
        # 显示所有检测到的人的关键点（半透明）
        for i, person_kpts in enumerate(self.all_detected_keypoints):
            if i != self.selected_person:  # 非选中的人显示为灰色
                for point in person_kpts:
                    if point is not None:
                        cv2.circle(display_image, point, 3, (128, 128, 128), -1)
        
        # 显示已标注的点和连接线
        connections = [
            (0, 1), (0, 2),  # 头部到肩膀
            (1, 3), (3, 5),  # 左臂
            (2, 4), (4, 6),  # 右臂
            (1, 7), (2, 8),  # 肩膀到臀部
            (7, 9), (9, 11),  # 左腿
            (8, 10), (10, 12)  # 右腿
        ]
        
        # 绘制连接线
        if self.keypoints:
            for i, j in connections:
                if i < len(self.keypoints) and j < len(self.keypoints) and \
                   self.keypoints[i] is not None and self.keypoints[j] is not None:
                    pt1 = self.keypoints[i]
                    pt2 = self.keypoints[j]
                    cv2.line(display_image, pt1, pt2, (0, 255, 255), 2)
            
            # 显示选中人物的关键点
            for i, point in enumerate(self.keypoints):
                if point is not None and i < len(self.keypoint_names):
                    x, y = point
                    cv2.circle(display_image, (x, y), 5, (0, 255, 0), -1)
                    # 使用中文字体显示关键点名称
                    display_image = self.put_chinese_text(display_image, 
                                                        self.keypoint_names[i], 
                                                        (x+5, y-25), 
                                                        self.small_font)
        
        # 显示提示信息
        if self.current_point < len(self.keypoint_names):
            display_image = self.put_chinese_text(
                display_image,
                f"请标注或调整: {self.keypoint_names[self.current_point]} (按空格跳过不可见点)", 
                (10, 30),
                color=(0, 0, 255)
            )
        else:
            display_image = self.put_chinese_text(
                display_image,
                "按's'保存并继续，'r'重置，'b'返回上一张",
                (10, 30),
                color=(0, 255, 0)
            )
        
        # 显示进度和人物选择提示
        progress = f"进度: {self.current_image_idx + 1}/{len(self.images)} | "
        progress += f"检测到 {len(self.all_detected_keypoints)} 人 | "
        progress += f"当前选择: 第 {self.selected_person + 1} 人 | "
        progress += "按数字键1-9选择不同人物"
        
        display_image = self.put_chinese_text(
            display_image,
            progress,
            (10, self.image_height - 40),
            self.small_font,
            color=(255, 255, 255)
        )
        
        # 显示已跳过的点
        skipped_points = [i for i, p in enumerate(self.keypoints) if p is None]
        if skipped_points:
            skipped_text = "已跳过的点: " + ", ".join([self.keypoint_names[i] for i in skipped_points])
            display_image = self.put_chinese_text(
                display_image,
                skipped_text,
                (10, 70),
                self.small_font,
                color=(0, 165, 255)  # 橙色
            )
        
        cv2.imshow(self.window_name, display_image)
    
    def save_annotations(self):
        annotations = {
            "image": self.images[self.current_image_idx],
            "keypoints": self.keypoints,
            "keypoint_names": self.keypoint_names,
            "image_size": [self.image_width, self.image_height]
        }
        
        output_dir = "annotations"
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            
        output_file = os.path.join(output_dir, 
                                 f"{self.images[self.current_image_idx].split('.')[0]}.json")
        
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(annotations, f, ensure_ascii=False, indent=2)
    
    def run(self):
        cv2.namedWindow(self.window_name)
        cv2.setMouseCallback(self.window_name, self.mouse_callback)
        
        # 初始化第一张图片
        self.keypoints = [None] * len(self.keypoint_names)
        self.load_image()
        
        while True:
            self.update_image()
            key = cv2.waitKey(1) & 0xFF
            
            # 数字键1-9用于选择不同的人
            if ord('1') <= key <= ord('9'):
                person_idx = key - ord('1')
                if person_idx < len(self.all_detected_keypoints):
                    self.select_person(person_idx)
            
            # 空格键跳过当前点
            elif key == ord(' '):
                if self.current_point < len(self.keypoint_names):
                    self.keypoints[self.current_point] = None
                    self.current_point += 1
                    print(f"已跳过: {self.keypoint_names[self.current_point-1]}")
            
            # 按's'保存当前标注
            elif key == ord('s'):
                # 检查是否所有点都已经标注或明确跳过
                if self.current_point >= len(self.keypoint_names):
                    self.save_annotations()
                    print(f"已保存标注: {self.images[self.current_image_idx]}")
                    
                    # 重置状态，准备下一张图片
                    self.current_image_idx += 1
                    if self.current_image_idx >= len(self.images):
                        break
                    self.keypoints = [None] * len(self.keypoint_names)
                    self.current_point = 0
                    self.current_image = None
                    self.all_detected_keypoints = []
                else:
                    print("请完成所有关键点的标注或跳过后再保存")
            
            # 按'q'退出
            elif key == ord('q'):
                break
            
            # 按'b'返回上一张图片
            elif key == ord('b') and self.current_image_idx > 0:
                self.current_image_idx -= 1
                self.keypoints = [None] * len(self.keypoint_names)
                self.current_point = 0
                self.current_image = None
                self.all_detected_keypoints = []
            
            # 按'r'重置当前图片的标注
            elif key == ord('r'):
                self.keypoints = [None] * len(self.keypoint_names)
                self.current_point = 0
                self.current_image = None
                self.all_detected_keypoints = []
        
        cv2.destroyAllWindows()

def main():
    # 首先处理训练集
    annotator = KeypointAnnotator("dataset/train")
    print("开始标注训练集...")
    print("操作说明：")
    print("- 程序会自动进行预标注")
    print("- 数字键1-9可以选择不同的人物")
    print("- 鼠标左键点击可以调整或标注关键点")
    print("- 空格键可以跳过不可见的关键点")
    print("- 's' 保存当前图片的标注并进入下一张")
    print("- 'r' 重置当前图片的标注")
    print("- 'b' 返回上一张图片")
    print("- 'q' 退出程序")
    annotator.run()

if __name__ == "__main__":
    main() 